import cv2
import numpy as np
import os
import cv2
from PIL import Image
from PIL import Image


def imgScale(img, newSize):
    # 等比例缩放至size大小的图片
    img.thumbnail((newSize, newSize))
    # 创建一个size大小的背景
    newImg = Image.new('RGB', (newSize, newSize), (0, 0, 0))

    # 获取图片的大小
    w, h = img.size

    # 将前景图片贴到背景图片上
    if w == newSize:
        newImg.paste(img, (0, int((newSize - h) / 2)))
    elif h == newSize:
        newImg.paste(img, (int((newSize - w) / 2), 0))
    else:
        newImg.paste(img, (int((newSize - w) / 2), int((newSize - h) / 2)))
    # newImg=np.array(newImg)[...,::-1]

    return newImg


def locate(img_src, img_mask, name='111', save_path=r'./data/'):
    """
    该函数通过cv2对img_mask进行边缘检测，获取车牌区域的边缘坐标(存储在contours中)和最小外接矩形4个端点坐标,
    再从车牌的边缘坐标中计算出和最小外接矩形4个端点最近的点即为平行四边形车牌的四个端点,从而实现车牌的定位和矫正
    :param img_src: 原始图片
    :param img_mask: 通过u_net预测得到的二值化图片，车牌区域呈现白色，背景区域为黑色
    :return: 定位且矫正后的车牌
    """
    ret, binary = cv2.threshold(img_mask, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    # cv2.imshow('1',_)
    # cv2.waitKey(0)
    # exit()
    # contours, hierarchy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # cv2.drawContours(img_mask, contours, -1, (255, 0, 0), 2)
    # cv2.imshow('', img_mask)
    # cv2.waitKey(0)
    # exit()

    abs = cv2.convertScaleAbs(binary, alpha=6, beta=0)
    # kerner = cv2.getStructuringElement(cv2.MORPH_RECT, (6, 6))
    # open = cv2.morphologyEx(abs, cv2.MORPH_OPEN, kerner, iterations=3)

    # close = cv2.morphologyEx(open, cv2.MORPH_CLOSE, kerner, iterations=1)
    gaussian = cv2.GaussianBlur(abs, (5, 5), 0)
    mask = cv2.Canny(gaussian, 50, 150)
    cv2.imshow('canny',mask)
    # cv2.imshow('2',img_src)
    # cv2.waitKey(0)
    # exit()
    # ret, _ = cv2.threshold(mask, 5, 100, cv2.THRESH_BINARY)
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # cv2.drawContours(img_src, contours, -1, (0, 0, 255), 2)
    # cv2.imshow('', img_src)
    # cv2.waitKey(0)
    # exit()

    if not len(contours):  # contours长度为0说明未检测到车牌
        print("未检测到车牌")
    else:
        flag = 0  # 默认flag为0，因为不一定有车牌区域
        for j in range(len(contours)):
            contour = contours if len(contours) == 1 else [contours[j]]

            for ii, cont in enumerate(contour):
                x, y, w, h = cv2.boundingRect(cont)  # 获取边界矩形
                img_cut_mask = img_mask[y:y + h, x:x + w]  # 将标签车牌区域截取出来
                # np.mean(img_cut_mask) >= 75
                if np.mean(img_cut_mask) >= 10 and w > 15 and h > 15:
                    rect = cv2.minAreaRect(cont)  # 针对坐标点获取带方向角的最小外接矩形，中心点坐标，宽高，旋转角度
                    box = cv2.boxPoints(rect).astype(np.int32)  # 获取最小外接矩形四个顶点坐标
                    cv2.drawContours(img_mask, contour, -1, (0, 0, 255), 2)
                    cv2.drawContours(img_mask, [box], 0, (0, 255, 0), 2)

                    cont = cont.reshape(-1, 2).tolist()
                    # 由于转换矩阵的两组坐标位置需要一一对应，因此需要将最小外接矩形的坐标进行排序，最终排序为[左上，左下，右上，右下]
                    box = sorted(box, key=lambda xy: xy[0])  # 先按照左右进行排序，分为左侧的坐标和右侧的坐标
                    box_left, box_right = box[:2], box[2:]  # 此时box的前2个是左侧的坐标，后2个是右侧的坐标
                    box_left = sorted(box_left, key=lambda x: x[1])  # 再按照上下即y进行排序，此时box_left中为左上和左下两个端点坐标
                    box_right = sorted(box_right, key=lambda x: x[1])  # 此时box_right中为右上和右下两个端点坐标
                    box = np.array(box_left + box_right)  # [左上，左下，右上，右下]

                    x0, y0 = box[0][0], box[0][1]  # 这里的4个坐标即为最小外接矩形的四个坐标
                    x1, y1 = box[1][0], box[1][1]
                    x2, y2 = box[2][0], box[2][1]
                    x3, y3 = box[3][0], box[3][1]

                    # 获取平行(或不规则)四边形的坐标
                    def point_to_line_distance(X, Y):
                        if x2 - x0:
                            k_up = (y2 - y0) / (x2 - x0)  # 斜率不为无穷大
                            d_up = np.abs(k_up * X - Y + y2 - k_up * x2) / (k_up ** 2 + 1) ** 0.5
                        else:  # 斜率无穷大
                            d_up = np.abs(X - x2)
                        if x1 - x3:
                            k_down = (y1 - y3) / (x1 - x3)  # 斜率不为无穷大
                            d_down = np.abs(k_down * X - Y + y1 - k_down * x1) / (k_down ** 2 + 1) ** 0.5
                        else:  # 斜率无穷大
                            d_down = np.abs(X - x1)
                        return d_up, d_down

                    d0, d1, d2, d3 = np.inf, np.inf, np.inf, np.inf
                    l0, l1, l2, l3 = (x0, y0), (x1, y1), (x2, y2), (x3, y3)
                    for each in cont:  # 计算cont中的坐标与矩形四个坐标的距离以及到上下两条直线的距离，对距离和进行权重的添加，成功选出四边形的4个顶点坐标
                        x, y = each[0], each[1]
                        dis0 = (x - x0) ** 2 + (y - y0) ** 2
                        dis1 = (x - x1) ** 2 + (y - y1) ** 2
                        dis2 = (x - x2) ** 2 + (y - y2) ** 2
                        dis3 = (x - x3) ** 2 + (y - y3) ** 2
                        # print(type(x),y)
                        # exit()
                        d_up, d_down = point_to_line_distance(float(x), float(y))
                        # weight = 0.975
                        weight = 0.8
                        if weight * d_up + (1 - weight) * dis0 < d0:
                            d0 = weight * d_up + (1 - weight) * dis0
                            l0 = [x, y]
                        if weight * d_down + (1 - weight) * dis1 < d1:
                            d1 = weight * d_down + (1 - weight) * dis1
                            l1 = [x, y]
                        if weight * d_up + (1 - weight) * dis2 < d2:
                            d2 = weight * d_up + (1 - weight) * dis2
                            l2 = [x, y]
                        if weight * d_down + (1 - weight) * dis3 < d3:
                            d3 = weight * d_down + (1 - weight) * dis3
                            l3 = [x, y]

                    p0 = np.array([l0, l1, l2, l3],
                                  dtype=np.float32)  # 左上角，左下角，右上角，右下角，形成的新box顺序需和原box中的顺序对应，以进行转换矩阵的形成
                    x1, y1 = l0[0], l0[1]
                    x2, y2 = l1[0], l1[1]
                    x3, y3 = l2[0], l2[1]
                    x4, y4 = l3[0], l3[1]

                    # print(p0)
                    # exit()
                    cv2.line(img_src,(x1,y1),(x2,y2),color=(255,0,0),thickness=2)
                    cv2.line(img_src,(x1,y1),(x3,y3),color=(255,0,0),thickness=2)
                    cv2.line(img_src,(x2,y2),(x4,y4),color=(255,0,0),thickness=2)
                    cv2.line(img_src,(x3,y3),(x4,y4),color=(255,0,0),thickness=2)
                    cv2.imshow('img_src',mask)
                    cv2.imshow('img_src',img_src)
                    # cv2.imshow('img_mask',mask)
                    cv2.waitKey(0)

                    # p1 = np.float32([[0, 0], [0, 512], [512, 0], [512, 512]])
                    # transform_mat = cv2.getPerspectiveTransform(p0, p1)  # 构成转换矩阵
                    # lic = cv2.warpPerspective(img_src, transform_mat, dsize=(512, 512))  # 进行透视变换车牌矫正
                    #
                    # if len(contour) == 1:  # 只有一个区域可以认为是车牌区域
                    #     flag += 1
                    #     print('saving ', save_path + name[0:7] + '.png')
                    #     cv2.imshow('lic', lic)
                    #     cv2.waitKey(0)
                    # cv2.imencode('.png', lic)[1].tofile(save_path + name[0:7] + '.png')

            if not flag:
                print("未检测到车牌区域或车牌区域过小")


if __name__ == '__main__':
    # img_path = r'./data/pictiure/river_level_0stream_cropped_12.jpg'
    img_path = r'./data/pics_at_night/stream0_1620743650.59416.jpg'
    mask_path = r'./data/test_imgs/stream0_1620743650.59416.jpg'
    save_path = './data/'  # 车牌图片保存路径
    img_mask = cv2.imread(mask_path, 0)
    cv2.imshow('',img_mask)
    # cv2.waitKey(0)
    src = Image.open(img_path)
    # src.show()
    # exit()
    # new_src = imgScale(src, 516)
    new_src = cv2.cvtColor(np.array(src), cv2.COLOR_RGB2BGR)
    locate(new_src, img_mask, '1111', save_path)
    exit()

    # mask_path = r'G:/00_deeplearning/01_code/data/box_progrem/u_net/split_image/'
    # img_path = r'G:/00_deeplearning/01_code/data/box_progrem/u_net/test_img/'
    # save_path = 'G:/00_deeplearning/01_code/data/box_progrem/box_cls_net/train/'  # 车牌图片保存路径
    # mask_list = os.listdir(mask_path)
    # rst_list = os.listdir(img_path)

    for i, img_mask in enumerate(mask_list):
        name = (img_mask.split('.')[0])
        img_src = name + '.jpg'

        src = os.path.join(img_path, img_src)
        mask = os.path.join(mask_path, img_mask)
        mask = cv2.imread(mask, 0)
        src = Image.open(src)

        new_src = imgScale(src, 516)
        new_src = cv2.cvtColor(np.array(new_src), cv2.COLOR_RGB2BGR)
        locate(new_src, mask, str(i), save_path)
